open Utils.Pervasive
open Utils.ListSpecial
open Utils.RandSelect
open Utils.Statistics
open Utils.StringSpecial

open Hashtbl
open Syntax
open Binding
open Core
open PrintExpr

open BlockPopulation.Gene 
open BlockPopulation 

open ProofAssembler
open Rules
open TestCases
open List
(*Private utils******************************)
let mprf str = prf str; read_line( )

type popSize = Max | MaxProportion of float | Fixed of int



type speciesInfo = {
	speciesId : int;
	ancesterID : int option;
	mutable fillProportion : float;
	mutable theoreticalMaxFill : int;
	mutable currentFill : int;
	mutable avgComplexity : float;
	mutable avgFitness : float;
	}
let emptyInfo id ancesterID = {
	speciesId = id;
	ancesterID = ancesterID;
	fillProportion = 0.;
	theoreticalMaxFill = 0;
	currentFill = 0;
	avgComplexity = 0.;
	avgFitness =0.;
	}	
type genoStatus = Dead of float | Alive of float | Unborn
type genotype = program

type population = {
	mutable info : speciesInfo;
	mutable contains_solution : bool;
	genotypes : (genotype, genoStatus) t 
	}
let empty_population id ancesterID = {
	info = emptyInfo id ancesterID;
	contains_solution =false;
	genotypes = create 200
	}
	
type species = proof

type ecoSetting = {
	maxProofWidth : int;
	aprNumberOfLiveGenotypes : int;
	aprGenePoolTotComp : int;
	selectSize : int;
	selectReap : int
	}
type ecosystem  = { 
	maxGenotypeScore : float;
	mutable solutions : term list;
	mutable speciesCtr : int;
	mutable generationCounter : int;
	mutable stats : string;
	settings : ecoSetting;
	goalType : ty; 
	mutable environment : test_case list;
	genes : gene_pool ref; 
	kingdom : (species,  population) t  
	} 
let get_species_id eco species = 
	let ctx = gpContext !(eco.genes) in
	let idOpt = 
		Hashtbl.fold 
		(fun spec pop acc -> 
			if (ProofAssembler.is_equiv ctx spec species) 
			then Some (pop.info.speciesId,pop.info.ancesterID)
			else acc)
		eco.kingdom None
	in match idOpt with 
	|None -> raise Not_found
	|Some id -> id	
let get_species_from_id eco id = 
	let speciesOpt = 
	Hashtbl.fold 
		(fun spec pop acc -> 
			if (pop.info.speciesId = id) 
			then Some spec 
			else acc) eco.kingdom None
	in match speciesOpt with 
		| None -> raise Not_found
		|Some species -> species	
let new_species_id eco = 
	let id = eco.speciesCtr in
	eco.speciesCtr <- eco.speciesCtr + 1;
	id	
let push eco stat = eco.stats <-   eco.stats ^ ("\n" ^ stat) 
let make_eco ?(debug=false) genes_ref goal test_cases settings  = 
	let new_eco = {
		maxGenotypeScore = float (List.length test_cases);
		solutions = [ ];
		stats = "";
		speciesCtr = 0;
		generationCounter = 0;
		settings = settings;
		goalType = goal;
		environment = test_cases;
		genes = genes_ref;
		kingdom = create 300} in
	new_eco	
let empty_eco genes = {
	maxGenotypeScore = 0.0;
	settings = {
		maxProofWidth = 10;
		aprNumberOfLiveGenotypes = 0;
		aprGenePoolTotComp = 0;
		selectSize = 0;
		selectReap = 0 
		};
	stats ="";
	solutions = [ ];
	speciesCtr = 0;
	generationCounter = 0 ;
	goalType = TyArr(TyUnit, TyUnit);
	environment = [empty_test_case];
	genes = genes;
	kingdom = create 0;
	}
let kill_all_genotypes eco = 
	ListLabels.iter  (keys eco.kingdom)
	~f:(fun species  -> 
		let (id , ancesterId)= get_species_id eco species in
		Hashtbl.replace (eco.kingdom) species (empty_population id ancesterId))
let rec string_of_genotype ?(separator="\n") genes =
	function 
		| In id -> string_of_int id
		| TypeFunc (prog, ty)-> (string_of_genotype ~separator genes prog) ^ "::" ^ (string_of_type (gpContext genes) ty Compact)
		| ObjFunc (prog1, prog2) -> (string_of_genotype ~separator genes prog1) ^  " ][ " ^ (string_of_genotype ~separator genes prog2)
let card_population ?(pred=(fun _-> true))  population = 
	Hashtbl.fold 
	(fun  genotype status acc ->
		if pred status then acc + 1 else acc) 
	population.genotypes 0
	
type eco_pop = Eco of ecosystem | Pop of population
let get_live pop = 
	match pop with 
	|Pop population -> 
		Hashtbl.fold (fun  genotype status lst ->
			match status with
			|Alive ev -> (ev , genotype) :: lst
			|_->lst) population.genotypes []
	|Eco eco ->	failwith "Comming"

	
	
let iter_population ?(pred=(fun _-> true)) ?(f=(fun _ _ -> ( )))   population = 
	Hashtbl.iter
	(fun genotype status -> if pred status then f genotype status)
	population
	
let card_ecosystem ?(pred=(fun _->true)) eco = 
	Hashtbl.fold 
	(fun _ population acc -> acc + (card_population ~pred population))  
	eco.kingdom 0
	 


let update eco species = 
	let genes = !(eco.genes) in
	let genotypes = (Hashtbl.find (eco.kingdom) species).genotypes in
	let genos = progs_of_proof genes species in 
	List.iter (fun genotype -> 
				if not (Hashtbl.mem genotypes genotype) then (
				Hashtbl.add genotypes genotype Unborn)) genos
		
		
let rec update_population ?(debug=false) eco species ancestorID =
	let deb ?(break=false) str = if debug then (if break then un(mprf str) else prf str) in
	let g_pop = !(eco.genes) in

	if  Hashtbl.mem eco.kingdom species 
	then (
		if is_complete species then (
			deb "\nSpecies add start";
			let population = 
				try Hashtbl.find eco.kingdom species 
				with Not_found -> assert(0=1); failwith "Shouldn't happen" in
		
			(try update eco species 				
				with _ -> failwith "Unable to construct population");
			let genotypes = population.genotypes in
			let liveGenos = get_live (Pop population) in
			let (id,aid) = get_species_id eco species in 
			let new_info = {
				speciesId = id;
				ancesterID = aid;
				theoreticalMaxFill = Hashtbl.length genotypes;
				currentFill = List.length liveGenos;
				fillProportion = 0.0;
				avgComplexity =  0.0;
				avgFitness = 0.0 } in 
			population.info <- new_info;
			deb "\nSpecies add end"))
	else (Hashtbl.add eco.kingdom species {
		info = emptyInfo (new_species_id eco) ancestorID;
		contains_solution = false;
		genotypes = Hashtbl.create 100};
		update_population eco species ancestorID)		
let add_species ?(debug=false) eco proof ancesterID = 
	let ctx = gpContext !(eco.genes) in
	assert (ty_eqv ctx (check_tree ctx proof) eco.goalType);
	update_population ~debug eco proof ancesterID;
	get_species_id eco proof

let candidateTypes genes gt = 
	get_types genes ~pred:(fun ty ->type_to (gpContext genes) gt ty)
let implyTypes goalType genes = 
	ListLabels.map (candidateTypes genes goalType)
	~f:(function 
		|ty when ty_eqv (gpContext genes) ty goalType -> ProgNode ty
		|ty -> Partial (goalType, Some (ProgNode ty))
		(*ImplyElim (ProgNode ty, Partial (ty1, None))*)
		|_->  assert (1=0);failwith "Not supported")
						
						
let rec mp proof eco depth goalType = 
	assert (depth >=1);
	let genes = !(eco.genes) in
	let ctx = gpContext genes in
	let proofType = check_tree ctx proof in
	if ty_eqv ctx proofType goalType 
	then(
		match depth with 
		|_ when depth = 1 -> 
			(match proof with 
			| Partial (tyGoal, Some proof) -> 
				ListLabels.map (mp proof eco 1 tyGoal)
				~f:(fun proof -> 
					if ty_eqv ctx tyGoal (check_tree ctx proof) 	
					then proof 
					else Partial(tyGoal, Some proof))
			| Partial (ty, None) -> implyTypes ty genes 
			| _->[proof])	
		|_->
			match proof with
			| ProgNode ty -> [ProgNode ty] 
			| Partial (tyGoal, _) -> mp proof eco 1 tyGoal
			| ImplyElim(proof1, proof2)->
				cartesianProduct 
				(fun proof1 proof2 -> ImplyElim(proof1,proof2))
				(mp proof1 eco (depth - 1) (check_tree ctx proof1)) 
				(mp proof2 eco (depth - 1) (check_tree ctx proof2)) 
			|_->  [ ])
	else(
		match proofType with 
		|TyArr(ty1, ty2) -> [ImplyElim (proof, (Partial (ty1,None)))]
		| _ -> assert (1=0);failwith "???"
	)		
let modifyProofs ?(proofs=[]) eco depth = 
	assert (depth >= 0);
	let genes = !(eco.genes) in
	match depth with
		|x when x = 0 -> 
			ListLabels.map (implyTypes eco.goalType genes)
			~f:(fun proof -> (None,proof))	
		|_ -> 
		let proofs = match proofs with
			|[ ] -> keys eco.kingdom
			|_->proofs in
		List.flatten (
			ListLabels.map proofs
			~f:(fun proof -> 
				let aid = fst (get_species_id eco proof) in
				ListLabels.map (mp proof eco depth eco.goalType)
				~f:(fun p -> (Some aid, p))))													
let speciesAccept ?(max=20) eco new_proofs =
	let genes  = !(eco.genes) in 
	let add proofLst = 
		ListLabels.iter proofLst
		~f:(fun (ancesterID, proof) -> 
			if not (ProofAssembler.exists (gpContext genes) (keys eco.kingdom) proof) 
			then (un(add_species eco proof ancesterID))) in	
	
	prf ("\n"^string_of_int (List.length new_proofs));
	let (new_complete_proofs, new_incomplete_proofs) = 
		ListLabels.partition new_proofs 
		~f:(fun (_,pr) -> 
			is_complete pr && 
			(not (ProofAssembler.exists (gpContext genes) (keys eco.kingdom) pr))) in
	
	prf ("\nAdding "^string_of_int (List.length new_complete_proofs) ^ " complete proofs");
	add new_complete_proofs;
	
	let lst =  
		if List.length new_incomplete_proofs < max 
		then new_incomplete_proofs
		else choose_random ~selectionSize:max new_incomplete_proofs in
	prf ("\nAdding "^string_of_int (List.length lst)^ " incomplete proof");
	add lst		 		
let makeProofs goOn eco  = 
	let depth = ref 0 in
	while goOn eco !depth do
		speciesAccept eco (modifyProofs eco !depth);
		depth := !depth + 1
	done	
	
let update_species_gene_use eco =
	let genes = !(eco.genes) in 
	reset_values genes;
	let rec update species = 
		match species  with 
		| ProgNode ty  -> 
			let catInfo, _ = tyFind genes ty in
			catInfo.species_use <- catInfo.species_use + 1
		| ForAllElim (proof, ty1)  -> update proof
		| ImplyElim (proof1, proof2) -> update proof1; update proof2
		| _->failwith "Invalid pair proof/program" in
	
	Hashtbl.iter (fun species _ -> update species) eco.kingdom
						
let evaluate eco species prog foundSolution =  
	let tests = eco.environment in
	let genes = !(eco.genes) in
	let ctx = gpContext genes in
	let te = term_of_program genes species prog in 
	let score = ListLabels.fold_left ~init:0.0 tests
	~f:(fun acc test -> if match_case te ctx test then acc +. 1.0 else acc) in
	match score with 
	|x when x = eco.maxGenotypeScore ->
		foundSolution := true;
		prf ("\nFound solution :" ^ string_of_term ctx te Compact ^ "\n");
		prf (string_of_program genes (species, prog));
		(Hashtbl.find eco.kingdom  species).contains_solution <- true;
		if not (List.exists (fun lstTe ->  te_eqv ctx lstTe te) eco.solutions) then  
		eco.solutions <- te :: eco.solutions;
		(Some te, score)
	|_ -> (None, score)

let evaluate_genotypes eco genotypeEvaluationNum = 
	assert(card_ecosystem eco > 0);			
	let prob = (float eco.settings.aprNumberOfLiveGenotypes) /. float (card_ecosystem eco) in
	let foundSolution = ref false in
	let sol = ref None in
	Hashtbl.iter (fun species population ->
		match !foundSolution with
		|true -> ( )
		|_-> Hashtbl.iter (fun prog status -> 
			match !foundSolution with |true -> ( )
			|_-> match status with 
			|Unborn -> 
				if Random.float 1.0 <= prob then(	
				prf "\nEvaluating genotype";			
				let (solOption, eval) = evaluate eco species prog foundSolution in
				Hashtbl.replace population.genotypes prog (Alive eval);
				sol := solOption;
				genotypeEvaluationNum := !genotypeEvaluationNum + 1)		
				
			| _ -> ( )) population.genotypes)
	eco.kingdom;
	!sol

let evaluate_species eco = 
	let es species population = 
		population.info.avgFitness <- 0.0;
		let num = ref 0.0 in
		let fit = ref 0.0 in
		Hashtbl.iter (fun _ status -> 
			match status with 
			| Alive ev -> 
				num := !num +. 1.0;
				fit := !fit +. ev
			|_->( )) population.genotypes;
		if !num > 0.0 then population.info.avgFitness <- !fit /. !num in
	Hashtbl.iter es eco.kingdom
	
let evaluate_genes eco =
	let genes = !(eco.genes) in
	let totFitAcc = ref 0.0 in
	reset_values genes;
		
	let rec update species prog ev = 
		match (species, prog) with 
		| ProgNode ty, In id -> 
			let catInfo, lst = tyFind genes ty in
			catInfo.species_use <- catInfo.species_use + 1;
			let alle = List.find (fun g -> alleId g = id) lst in
				alleIncUseInPop alle;
				alleSetFitness alle ((alleFitness alle) +. ev);
				if (alleScoreOfBestCarrier alle) < ev then alleSetScoreBestCarr alle ev
		| ForAllElim (proof, ty1), TypeFunc (prog, ty2) -> 
			update proof prog ev
		| (ImplyElim (proof1, proof2)) , (ObjFunc (prog1, prog2)) ->
			update proof1 prog1 ev; update proof2 prog2 ev
		| _->failwith "Invalid pair proof/program" in
	Hashtbl.iter (fun species population ->
		Hashtbl.iter (fun prog -> function 
			| Alive ev -> update species prog ev;
			|_->( )) population.genotypes) eco.kingdom;
	
	 	
	 gpIter (fun _ cinfo geneLst ->
		match cinfo.species_use with
		|x when x > 0 -> List.iter alleNormalizeFitness geneLst 
		|_->( )) genes


let fill_gene_pool constraintFunc genes maxComp =
		let compGenes = comp genes in
		match maxComp - compGenes with
		| x when x > 0 ->
				growToComplexity 
				~growUnit:(Complexity x)
				genes constraintFunc
				~onSuccess:(fun bpset _ -> prf "\nsuccess")
				~onFailure:(fun str -> prf "\nfailure")
		|_ -> prf "\n Done building gene pool"	
			

let evaluate_ecosystem eco genotypeEvaluationNum = 
		let sol = evaluate_genotypes eco genotypeEvaluationNum in
		evaluate_species eco;
		evaluate_genes eco;
		eco.generationCounter <- eco.generationCounter + 1;
		sol	
		
let rec init_generation ?(debug=false) eco constraintFunc genotypeEvaluationNum = 
	push eco ("----\nGeneration " ^ (string_of_int eco.generationCounter));
	let genes = !(eco.genes) in
	let prog_num ( ) = ProofAssembler.prog_num genes (keys eco.kingdom) in
	let ctr = ref 20 in
	let cond1( ) = !ctr > 0	 in
	let cond2( ) =	(prog_num ( )) <=  eco.settings.aprNumberOfLiveGenotypes in
	let cond3( ) = (Hashtbl.length eco.kingdom) < 15 in
		while ((cond1()) && (cond2())) && (cond3())  do
			fill_gene_pool constraintFunc genes eco.settings.aprGenePoolTotComp;
			
			let goOn = (fun eco depth -> if depth >= 6 then false else true) in
			makeProofs goOn eco;
	
			ListLabels.iter (keys eco.kingdom)
			~f: (fun species -> 
				if not (is_complete(species)) 
				then Hashtbl.remove eco.kingdom species);
		
			List.iter (fun species -> update eco species) (keys eco.kingdom);
			update_species_gene_use eco;
			gpRemoveUnusableGenes genes;
			ctr := !ctr - 1
		done;  
		push eco ("\t"^ (string_of_int (List.length (keys eco.kingdom))) ^ " species");
		push eco ("\t"^ string_of_int (card_ecosystem eco) ^ " possible genotypes");
		evaluate_ecosystem eco genotypeEvaluationNum  


(*and evalGenes eco constraintFunc = 
	let mean genes  = 
		let geneLst = get_genes genes in
		if List.length geneLst > 0 
		then (Utils.Statistics.mean geneLst (fun g->g.selectionProb))
		else failwith "There are no genes in the gene pool at mean"  in
	let stats genes = push eco 
		("\nCurrent gene pool size is " ^ (string_of_int (card genes)) ^ 
		"\ngene pool mean fitness is = " ^ (string_of_float (mean genes))) in
	let genes = !(eco.genes) in
	let totFitAcc = ref 0.0 in
	reset_values genes;
		
	let rec update species prog ev = 
		match (species, prog) with 
		| ProgNode ty, In id -> 
			let catInfo, lst = tyFind genes ty in
			catInfo.species_use <- catInfo.species_use + 1;
			let gene = List.find (fun g -> g.id = id) lst in
			gene.useInPopulation <- gene.useInPopulation +1;
			gene.geneFitness <- gene.geneFitness  +. ev;
			if gene.scoreOfBestCarrier < ev then gene.scoreOfBestCarrier <- ev
		| ForAllElim (proof, ty1), TypeFunc (prog, ty2) -> 
			update proof prog ev
		| (ImplyElim (proof1, proof2)) , (ObjFunc (prog1, prog2)) ->
			update proof1 prog1 ev; update proof2 prog2 ev
		| _->failwith "Invalid pair proof/program" in
	
	Hashtbl.iter (fun species population ->
		Hashtbl.iter (fun prog -> function 
			| Alive ev -> totFitAcc := !totFitAcc +. ev; update species prog ev;
			|_->( )) population.genotypes) eco.kingdom;
			
	 Hashtbl.iter (fun ty (cinfo, geneLst) -> 	
			if cinfo.species_use = 0 then rem_type genes ty else
			ListLabels.iter geneLst
			~f:(function gene -> 
				gene.geneFitness <- gene.geneFitness ;
				if gene.useInPopulation > 0 then
				gene.selectionProb <- gene.geneFitness /. float gene.useInPopulation)) 
			genes.g;
	
	stats genes;
	let mean = mean genes in
	Hashtbl.iter (fun ty (cinfo, geneLst) -> 	
			let id = ref 0 in
			let gLst = ListLabels.fold_left ~init:[ ] geneLst 
			~f:( fun lst gene -> 
				if gene.selectionProb > (Random.float mean *. 3.0) 
				then (
					gene.id <- !id;
					id := !id + 1;
					gene::lst) 
				else lst) in
			if (List.length gLst > 0) 
			then (cinfo.next_id <- !id;
			replace_type genes ty (Some cinfo) (Some gLst))
			else rem_type genes ty) genes.g;	
	prf "\nFinished evaluating gene pool";
	if card genes > 0 then stats genes
let select_for_reap eco	= 
	(*let size = ref (eco.settings.selectReap) in*)
	let cCh = function
		|Alive x -> Dead x
		|x -> x in

	let max = float (List.length eco.environment) in
	let pred = function 
		|Alive x (*when !size > 0 *) -> 
			x  <  Random.float 1.0 *. max *. 1.5
		|_->false in
	
	
	Hashtbl.iter (fun species population ->
		iter_population ~pred 
			~f:(fun genotype status -> 
				Hashtbl.replace population.genotypes 
				genotype (cCh status)) population.genotypes(*;
				size := !size -1*))
	eco.kingdom		   *)
	
 

		